package com.tsing.websocket.server;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;

public class SocketHandler implements WebSocketHandler {

    private final Queue<WebSocketSession> sessions = new ConcurrentLinkedQueue<>();

    private final Map<String, MessageDTO> lastMessages = new ConcurrentHashMap<>();

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        sessions.add(session);
        Mono<Void> sendLastMessagesToNewClient = sendLastMessagesToClient(session);
        return session.receive()
                .map(WebSocketMessage::getPayloadAsText)
                .map(this::toObject)
                .doOnNext(messageDTO -> {
                    lastMessages.put(session.getId(), messageDTO);
                    broadcastMessage(session, toJsonStr(new ResultDTO<>(SocketTypeEnum.SUCCESS.getCode(), messageDTO)));
                })
                .then()
                .doFinally(signalType -> {
                    MessageDTO messageDTO = lastMessages.get(session.getId());
                    if (Objects.nonNull(messageDTO)) {
                        broadcastMessage(session, toJsonStr(new ResultDTO<>(SocketTypeEnum.DISCONNECTED.getCode(), messageDTO.id())));
                    }
                    lastMessages.remove(session.getId());
                    sessions.remove(session);
                })
                .and(sendLastMessagesToNewClient);
    }

    private void broadcastMessage(WebSocketSession session, String message) {
        sessions.forEach(s -> {
            if (!s.getId().equals(session.getId())) {
                s.send(Mono.just(session.textMessage(message)))
                        .onErrorResume(e -> {
                            System.err.println("Error sending message to session " + s.getId() + ": " + e.getMessage());
                            return Mono.empty();
                        }).subscribe();
            }
        });
    }

    private Mono<Void> sendLastMessagesToClient(WebSocketSession newClient) {
        return Mono.defer(() -> newClient.send(Mono.just(newClient.textMessage(toJsonStr(new ResultDTO<>(SocketTypeEnum.CONNECTED.getCode(), lastMessages.values()))))));
    }

    private String toJsonStr(Object object) {
        try {
            return OBJECT_MAPPER.writeValueAsString(object);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private MessageDTO toObject(String jsonStr) {
        try {
            return OBJECT_MAPPER.readValue(jsonStr, MessageDTO.class);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
